

class CommandGraph:
    """
    Class for storing and parsing a declarative form of a OpenCL command graph
    
    The graph is described with a simple JSON syntax. See examples.
    
    
    """
    def __init__(self, clist, mem_objs, kernels):
        self.clist = clist
        self.mem_objs = mem_objs
        self.kernels = kernels
    
        self.symbols = {}
        self.symbols.update(mem_objs)
        self.symbols.update(kernels)


    def generate_DOT_string(self):
        # Generate a description of the command graph
        # in the DOT format (usable for Graphviz tools)
        dot_string = ["digraph opencl_command_graph {"]
        
        # Legend
        dot_string.append("subgraph cluster_legend {")
        dot_string.append("graph [rankdir=TB]")
        dot_string.append("node [shape=box] cnode1 [label=\"command 1\"]")
        dot_string.append("node [shape=box] cnode2 [label=\"command 2\"]")
        dot_string.append("node [shape=circle] mnode [label=\"memory\\nobject\"]")
        dot_string.append("cnode1 -> cnode2 [label=\"command\\ndependency\", labelfloat=true]")
        dot_string.append("cnode2 -> mnode [style=dotted, dir=both, label=\"memory\\naccess\", labelfloat=true]")
        dot_string.append("label = Legend")
        dot_string.append("}")
        
        # Command nodes
        command_nodes = ["node [shape=box]"]
        for command in self.clist:
            command_name = command[0]
            command_type = command[1]
            command_nodes.append("%s [label=\"%s : %s\"]"%(command_name, command_name, command_type))
        dot_string.append(";".join(command_nodes))
        
        # Memory object nodes
        mem_nodes = ["node [shape=circle]"]
        for mem_name in self.mem_objs:
            mem_nodes.append(mem_name)
        dot_string.append(";".join(mem_nodes))

        # Command dependencies        
        for command in self.clist:
            command_name = command[0]
            for dep_name in command[2].get("wait_for",[]):
                dot_string.append("    %s -> %s"%(dep_name, command_name))

        # Memory access for commands
        for command in self.clist:
            command_name = command[0]
            command_type = command[1]
            
            if(command_type=="nd_range_kernel"):
                for arg in command[2].get("arguments",[]):
                    if arg in self.symbols:
                        dot_string.append("    %s -> %s [style=dotted, dir=both]"%(arg,command_name))
            
            if(command_type=="copy"):
                dot_string.append("    %s -> %s [style=dotted]"%(command[2]["src"],command_name))
                dot_string.append("    %s -> %s [style=dotted]"%(command_name,command[2]["dst"]))
            
        dot_string.append("}")
        
        dot_string = "\n".join(dot_string)

        return dot_string
    
    def add_symbol(self, symbol, value):
        self.symbols[symbol] = value

    def integer_convert(self, argument):
        if(isinstance(argument, (int, long))):
            argument = np.int32(argument)
        return argument
        
    def integer_array_convert(self, argument_array):
        return [ self.integer_convert(arg) for arg in argument_array]
        
    def symbol_substitute(self, symbol):
        val = self.symbols.get(symbol, symbol)
        return val
        
    def symbol_array_substitute(self, symbol_array):
        values = [ self.symbol_substitute(symbol) for symbol in symbol_array]
        return values

    def parse_command(self, queue, command):
        command_type = command[1]
        command_name = command[0]
        command_params = command[2]
        handler_fun = getattr(self, "handle_%s" % command_type)
        cevent = handler_fun(queue, command_params)
        self.add_symbol(command_name, cevent)        

    def parse_command_list(self, queue):
        for command in self.clist:
            self.parse_command(queue, command)        
        
    def handle_copy(self, queue, params):
            
        wait_for = params.get("wait_for",[])
        wait_for = self.symbol_array_substitute(wait_for)
        
        dst = self.symbol_substitute(params["dst"])
        src = self.symbol_substitute(params["src"])
        
        event = pocl.enqueue_copy( queue, dst, src, wait_for = wait_for)
        
        return event

    def handle_nd_range_kernel(self, queue, params):
                
        kernel = self.symbol_substitute(params["kernel_name"])
        
        wait_for = params.get("wait_for",[])
        wait_for = self.symbol_array_substitute(wait_for)
    
        global_work_size = params["global_work_size"]
        local_work_size =  params["local_work_size"]
        
        arguments = params.get("arguments", [])
        arguments = self.symbol_array_substitute(arguments)
        arguments = self.integer_array_convert(arguments)
    
        kernel.set_args(*arguments)

        event = pocl.enqueue_nd_range_kernel(queue, kernel, global_work_size, local_work_size, global_work_offset=None, g_times_l=False, wait_for=wait_for)
              
        return event
    
    def handle_marker(self, queue, params):
        event = pocl.enqueue_marker(queue)
        return event
        

def enqueue_command_graph(queue, command_graph):
    
    command_graph.parse_command_list(queue)


if __name__ == "__main__":
    import pyopencl as pocl
    import numpy as np
    import numpy.linalg as la

    mflags = pocl.mem_flags
    ctx = pocl.Context(dev_type=pocl.device_type.CPU)
    queue = pocl.CommandQueue(ctx)
    
    A_host = np.random.rand(640*480).reshape([640, 480]).astype(np.float32)
    D_host = np.empty_like(A_host).astype(np.float32)
    D_reference = np.array(5).astype(np.float32)*np.sqrt(A_host) + (np.sqrt(A_host) + np.array(4.0).astype(np.float32))

    
    mem_objs = { "A_host": A_host,
                 "A_in":   pocl.Buffer(ctx, mflags.READ_ONLY, A_host.nbytes),
                 "A_out":  pocl.Buffer(ctx, mflags.READ_WRITE, A_host.nbytes),
                 "B_out":  pocl.Buffer(ctx, mflags.READ_WRITE, A_host.nbytes),
                 "C_out":  pocl.Buffer(ctx, mflags.READ_WRITE, A_host.nbytes),
                 "D_out":  pocl.Buffer(ctx, mflags.READ_WRITE, A_host.nbytes),
                 "D_host": D_host }
    
    program_text = """
                   __kernel void A_kernel(__global const float* A_in, __global float* A_out, int stride)
                   {
                       int gid0 = get_global_id(0);
                       int gid1 = get_global_id(1);
                       
                       A_out[gid1*stride + gid0] = sqrt(A_in[gid1*stride + gid0]);
                   }
                   
                   __kernel void B_kernel(__global const float* B_in, __global float* B_out, int stride)
                   {
                       int gid0 = get_global_id(0);
                       int gid1 = get_global_id(1);
                       
                       B_out[gid1*stride + gid0] = 5*B_in[gid1*stride + gid0];
                   }
                   
                   __kernel void C_kernel(__global const float* C_in, __global float* C_out, int stride)
                   {
                       int gid0 = get_global_id(0);
                       int gid1 = get_global_id(1);
                       
                       C_out[gid1*stride + gid0] = C_in[gid1*stride + gid0] + 4.0f;
                   }
                   
                   __kernel void D_kernel(__global const float* D_in1, __global const float* D_in2, __global float* D_out, int stride)
                   {
                       int gid0 = get_global_id(0);
                       int gid1 = get_global_id(1);
                       
                       D_out[gid1*stride + gid0] = D_in1[gid1*stride + gid0] + D_in2[gid1*stride + gid0];
                   }                                                   
                   """
    
    program = pocl.Program(ctx, program_text).build()
    
    kernels = {"A_kernel": program.A_kernel,
               "B_kernel": program.B_kernel,
               "C_kernel": program.C_kernel,
               "D_kernel": program.D_kernel}
    
    command_list = [ ["A_copy", "copy", 
                      {"dst": "A_in",
                       "src": "A_host"}],
                     ["A", "nd_range_kernel",
                      {"kernel_name": "A_kernel",
                       "arguments": ["A_in","A_out", A_host.shape[0]],
                       "global_work_size": list(A_host.shape),
                       "local_work_size": None,
                       "wait_for": ["A_copy"]}],
                     ["B", "nd_range_kernel",
                      {"kernel_name": "B_kernel",
                       "arguments": ["A_out","B_out", A_host.shape[0]],
                       "global_work_size": list(A_host.shape),
                       "local_work_size": None,
                       "wait_for": ["A"]}],
                     ["C", "nd_range_kernel",
                      {"kernel_name": "C_kernel",
                       "arguments": ["A_out","C_out", A_host.shape[0]],
                       "global_work_size": list(A_host.shape),
                       "local_work_size": None,
                       "wait_for": ["A"]}],
                     ["M", "marker", {}],
                     ["D", "nd_range_kernel",
                      {"kernel_name": "D_kernel",
                       "arguments": ["B_out","C_out","D_out", A_host.shape[0]],
                       "global_work_size": list(A_host.shape),
                       "local_work_size": None,
                       "wait_for": ["B","C","M"]}],
                     ["D_copy", "copy",
                      {"dst": "D_host",
                       "src": "D_out",
                       "wait_for":["D"]}]]
        
    command_graph = CommandGraph(command_list, mem_objs, kernels)
    enqueue_command_graph(queue, command_graph)
    
    print command_graph.generate_DOT_string()
    
    print la.norm(D_host - D_reference )
    print np.max(np.abs(D_host-D_reference))
